//
// Created by User on 24-10-9.
//
#include "general_muxer.h"
#include "packet_provider.hpp"
#include "media_packet.h"

#include <glog/logging.h>
#include <algorithm>
#include <filesystem>

extern "C"
{
#include <libavformat/avformat.h>
}

namespace muxer
{
    using namespace std;

    general_muxer::general_muxer(const std::string& path, const std::vector<interfaces::packet_provider*>& providers)
        : m_providers(providers), pts_vec(this->m_providers.size(), 0),
          pkt(av_packet_alloc()), marker(pts_vec.size(), true)
    {
        try
        {
            this->init_muxer(path);
        }
        catch (error::muxer_error& error)
        {
            this->clean_muxer(error.get_type());
            throw;
        }
        catch (...)
        {
            this->clean_muxer(error::muxer_error::type::UNKNOWN_ERROR);
            throw;
        }
    }

    general_muxer::~general_muxer()
    {
        this->clean_muxer(error::muxer_error::type::UNKNOWN_ERROR);
    }

    bool general_muxer::start_mux()
    {
        //写入文件头
        int ret = avformat_write_header(this->m_fmt_ctx, nullptr);
        if (ret < 0)
        {
            throw error::muxer_error(error::muxer_error::type::WRITE_HEADER_ERROR, ret);
        }
        //写入数据包
        this->write_pkt();
        if (this->pause_flag)
        {
            LOG(INFO) << "Pausing muxer...";
            return false;
        }
        //写入文件尾
        ret = av_write_trailer(this->m_fmt_ctx);
        if (ret < 0)
        {
            throw error::muxer_error(error::muxer_error::type::WRITE_TRAILER_ERROR, ret);
        }
        return true;
    }

    int general_muxer::current_frames() const
    {
        return this->m_curr_frames;
    }

    void general_muxer::init_muxer(const std::string& path)
    {
        if (!filesystem::exists(filesystem::path(path).parent_path()))
        {
            throw error::muxer_error(error::muxer_error::type::PATH_ERROR, "invalid path");
        }

        // 分配format ctx
        int ret = avformat_alloc_output_context2(&this->m_fmt_ctx, nullptr, nullptr,
                                                 path.c_str());
        if (ret < 0)
        {
            throw error::muxer_error(error::muxer_error::type::FMT_CTX_ALLOC_ERROR, ret);
        }

        // 添加各个流
        for (const auto provider: this->m_providers)
        {
            auto stream = avformat_new_stream(this->m_fmt_ctx, nullptr);
            if (!stream)
            {
                throw error::muxer_error(error::muxer_error::type::STREAM_ADD_ERROR, "can not add new stream");
            }
            ret = avcodec_parameters_copy(stream->codecpar, provider->get_codec());
            if (ret < 0)
            {
                throw error::muxer_error(error::muxer_error::type::PARAMETERS_SET_ERROR, ret);
            }
            this->m_streams.push_back(stream);
            LOG(INFO) << "new stream : " << stream->index;
        }
        // 打印输出格式详情
        av_dump_format(this->m_fmt_ctx, 0, path.c_str(), 1);

        //打开输出文件
        if (!(m_fmt_ctx->oformat->flags & AVFMT_NOFILE))
        {
            ret = avio_open(&m_fmt_ctx->pb, path.c_str(), AVIO_FLAG_WRITE);
            if (ret < 0)
            {
                throw error::muxer_error(error::muxer_error::type::FILE_OPEN_ERROR, ret);
            }
        }
        else
        {
            throw error::muxer_error(error::muxer_error::type::FILE_OPEN_ERROR, "can not found output file!");
        }
    }

    void general_muxer::clean_muxer(error::muxer_error::type _type)
    {
        switch (_type)
        {
            case error::muxer_error::type::UNKNOWN_ERROR: avio_closep(&this->m_fmt_ctx->pb);
            case error::muxer_error::type::FILE_OPEN_ERROR:
            case error::muxer_error::type::PARAMETERS_SET_ERROR:
            case error::muxer_error::type::STREAM_ADD_ERROR: avformat_free_context(this->m_fmt_ctx);
            case error::muxer_error::type::FMT_CTX_ALLOC_ERROR:
            case error::muxer_error::type::PATH_ERROR: ;
            default: ;
        }
    }

    void general_muxer::write_pkt()
    {
        //        vector pts_vec(this->m_providers.size(), 0);
        //        AVStream* output_stream{};
        //        AVRational input_timebase{};
        //        utils::media_packet pkt(av_packet_alloc());
        //        int end_counter = 0;
        //        vector marker(pts_vec.size(), true);
        while (end_counter < pts_vec.size() && !this->pause_flag)
        {
            int min_index = this->min_pts(pts_vec, marker);
            pkt = this->m_providers[min_index]->get_packet();
            if (!pkt)
            {
                end_counter++;
                marker[min_index] = false;
                continue;
            }
            input_timebase = this->m_providers[min_index]->get_timebase();
            output_stream = this->m_fmt_ctx->streams[min_index];
            pts_vec[min_index] = pkt.get_packet()->pts;
            pkt.get_packet()->stream_index = min_index;
            this->write_single_pkt(pkt, input_timebase, output_stream);
        }
    }

    int general_muxer::min_pts(const std::vector<int>& pts_vec, const vector<bool>& marker)
    {
        vector<double> vec;
        transform(this->m_providers.begin(), this->m_providers.end(),
                  back_inserter(vec), [](const interfaces::packet_provider* p)
                  {
                      return av_q2d(p->get_timebase());
                  });
        for (int i = 0; i < pts_vec.size(); i++)
        {
            vec[i] = vec[i] * pts_vec[i];
        }
        int ret = -1;
        double current = numeric_limits<int>::max();
        for (int i = 0; i < vec.size(); i++)
        {
            if (marker[i])
            {
                if (current > vec[i])
                {
                    ret = i;
                    current = vec[i];
                }
            }
        }
        return ret;
    }

    void general_muxer::write_single_pkt(utils::media_packet& pkt, AVRational input_timebase, AVStream* output_stream)
    {
        auto* temp_pkt = pkt.get_packet();
        temp_pkt->pts = av_rescale_q_rnd(
            temp_pkt->pts, input_timebase, output_stream->time_base,
            static_cast<AVRounding>(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
        temp_pkt->dts = av_rescale_q_rnd(
            temp_pkt->dts, input_timebase, output_stream->time_base,
            static_cast<AVRounding>(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
        temp_pkt->duration = av_rescale_q(temp_pkt->duration, input_timebase, output_stream->time_base);
        LOG(INFO) << "stream : " << pkt.get_packet()->stream_index << "\t"
                << temp_pkt->dts << "\t" << temp_pkt->pts << "\t" << temp_pkt->duration;
        int ret = av_interleaved_write_frame(this->m_fmt_ctx, pkt.get_packet());
        if (ret < 0)
        {
            throw error::muxer_error(error::muxer_error::type::WRITE_PACKET_ERROR, ret);
        }
        ++this->m_curr_frames;
    }

    void general_muxer::pause_mux()
    {
        this->pause_flag.store(true);
    }

    bool general_muxer::continue_mux()
    {
        this->pause_flag.store(false);
        this->write_pkt();
        if (this->pause_flag)
        {
            return false;
        }
        //写入文件尾
        if (const int ret = av_write_trailer(this->m_fmt_ctx); ret < 0)
        {
            throw error::muxer_error(error::muxer_error::type::WRITE_TRAILER_ERROR, ret);
        }
        return true;
    }
}
